import random
import numpy as np
from const import *  # 确保已定义 NEIGHBOR_LEN, HIS_LEN, HIDDEN


def graph_embedding_expansion(Otraining, usernei, alluserembs):
    """
    简化后的本地图拓展函数，去除了加密和第三方交互，直接使用 alluserembs 进行嵌入扩展。
    Otraining: 原始训练矩阵（未使用，可按需保留或移除）
    usernei: 每个用户的邻居 ID 列表，形如 List[List[int]]
    alluserembs: 一个字典或数组，key 为用户/物品 ID，value 为对应的 embedding 向量
    返回: shape 为 (num_users, HIS_LEN, NEIGHBOR_LEN, HIDDEN) 的 numpy 数组
    """
    user_neighbor_emb = []
    # 遍历每个用户
    for user_items in usernei:
        all_neighbor_embs = []
        # 对每个邻居 ID 获取 embedding 列表
        for item in user_items:
            # 直接使用该 ID 对应的 embedding
            emb = alluserembs[item]
            # 构建固定长度的邻居列表
            neighbor_embs = [emb]
            # 随机打乱（如需）并截断/填充到 NEIGHBOR_LEN
            random.shuffle(neighbor_embs)
            neighbor_embs = neighbor_embs[:NEIGHBOR_LEN]
            neighbor_embs += [[0.] * HIDDEN] * (NEIGHBOR_LEN - len(neighbor_embs))
            all_neighbor_embs.append(neighbor_embs)
        # 如果用户历史长度不足，填充到 HIS_LEN
        while len(all_neighbor_embs) < HIS_LEN:
            all_neighbor_embs.append([[0.] * HIDDEN] * NEIGHBOR_LEN)
        # 转为 numpy
        all_neighbor_embs = np.array(all_neighbor_embs, dtype='float32')
        user_neighbor_emb.append(all_neighbor_embs)

    # 对所有用户按最长历史长度进行统一 padding
    max_len = max(user.shape[0] for user in user_neighbor_emb)
    padded_user_neighbor_emb = []
    for user in user_neighbor_emb:
        pad_len = max_len - user.shape[0]
        if pad_len > 0:
            padding = np.zeros((pad_len, NEIGHBOR_LEN, HIDDEN), dtype='float32')
            padded_user = np.concatenate([user, padding], axis=0)
        else:
            padded_user = user
        padded_user_neighbor_emb.append(padded_user)

    return np.array(padded_user_neighbor_emb, dtype='float32')
